import numpy as np

class AntColony(object):
    def __init__(self, distance, n_ants, n_best, n_iterations, decay, alpha = 1, beta = 1):
        '''
        距离方阵
        每次迭代的蚂蚁的数量
        积累信息素的最佳蚂蚁数量
        迭代次数
        衰减系数
        用于控制信息素影响的参数,[0,1]
        用于控制可取性的参数,[0,1]
        '''
        self.distance = distance
        self.pheromone = np.ones(self.distance.shape) #tau
        self.all_inds = range(len(distance))
        self.n_ants = n_ants
        self.n_best = n_best
        self.n_iterations = n_iterations
        self.decay = decay
        self.alpha = alpha
        self.beta = beta
    #构成解
    def pick_move(self, pheromone, dist, visited):
        pheromone = np.copy(pheromone)
        #忽略已经去过的城市，因为回不去了
        pheromone[list(visited)] = 0
        #计算转移概率
        eta = 1/dist
        row = (pheromone ** self.alpha) *( eta**self.beta)
        norm_row = row/row.sum()
        move = np.random.choice(self.all_inds, 1, p = norm_row)[0]
        return move
    #信息素更新
    def spread_pheronome(self, all_paths,n_best):
        #优先考虑城市间最短的路线
        sorted_paths = sorted(all_paths,key = lambda x:x[1])

        for path, dist in sorted_paths[:n_best]:
            for move in path:
                self.pheromone[move] += 1.0/self.distance[move]

    def gen_all_paths(self):
        #生成蚂群路径和距离
        all_paths = []
        for i in range(self.n_ants):
            path = self.gen_path(0)
            all_paths.append((path, self.gen_path_dist(path)))
        return all_paths
    def gen_path(self, start):
        #生成单个蚂蚁的路径
        path = []
        visited = set()
        visited.add(start)
        prev = start
        #生成路径，起始点为start,中间点为prev
        #从一个城市走到另一个城市，至到访问所有的城市
        for i in range(len(self.distance) - 1):
            move = self.pick_move(self.pheromone[prev], self.distance[prev], visited)
            path.append((prev,move))
            prev = move
            visited.add(move)
        path.append((prev,start))
        return path
    def gen_path_dist(self, path):
        total_dist = 0
        for ele in path:
            total_dist += self.distance[ele]
        return total_dist
    def run(self):
        all_time_shortest_path = ("placeholder", np.inf)#(路径，距离)
        for i in range(self.n_iterations):
            all_paths = self.gen_all_paths()
            self.spread_pheronome(all_paths, self.n_best)
            #根据距离将路径排序
            shortest_path = min(all_paths,key = lambda x:x[1])
            print(shortest_path)
            #距离比较
            if shortest_path[1] < all_time_shortest_path[1]:
                all_time_shorest_path = shortest_path
            #信息素衰减，这样旧的信息素就不会迷惑下一代的蚂蚁
            #信息素矩阵乘以衰减率
            self.pheromone = self.pheromone * self.decay
        return all_time_shortest_path

if __name__ == "__main__":
    d = np.array([
        [np.inf, 20, 12, 3, 14],
        [10,np.inf,13, 15,  8],
        [12, 7, np.inf,9, 14],
        [1, 15, 9, np.inf, 16],
        [4, 8, 5, 16, np.inf]])
    ant_colony = AntColony(d, 100, 20, 10, 0.95, alpha = 1, beta = 2)
    ant_colony.run()


